#!/usr/bin/env python3
import open3d as o3d
import numpy as np
import argparse
import os
from pathlib import Path


def parse_path_file(path_file, is_outer=False):
    """解析路径文件，返回 PathPoint 列表：(x,y,z,rx,ry,rz,region)
    支持空格分隔和逗号分隔两种格式。"""
    points = []
    with open(path_file, 'r', encoding='utf-8') as f:
        for ln in f:
            ln = ln.strip()
            if not ln or ln.startswith('#'):
                continue
            
            # 检测是逗号分隔还是空格分隔
            if ',' in ln:
                parts = ln.split(',')
            else:
                parts = ln.split()
                
            # 内部文件: X Y Z RX RY RZ
            # 外部文件: X Y Z RX RY RZ Intensity(Region)
            try:
                if is_outer and len(parts) >= 7:
                    x, y, z, rx, ry, rz, region = parts[:7]
                    points.append((float(x), float(y), float(z), float(rx), float(ry), float(rz), int(float(region))))
                elif len(parts) >= 6:
                    x, y, z, rx, ry, rz = parts[:6]
                    points.append((float(x), float(y), float(z), float(rx), float(ry), float(rz), 0))
            except ValueError:
                # 忽略无法解析的行
                continue
    return points


def make_line_set(points):
    """根据点创建线集合，用灰色线连接相邻点"""
    if len(points) < 2:
        return None
    pts = np.asarray([[p[0], p[1], p[2]] for p in points], dtype=float)
    lines = [[i, i+1] for i in range(len(pts)-1)]
    colors = [[0.6, 0.6, 0.6] for _ in lines]
    ls = o3d.geometry.LineSet()
    ls.points = o3d.utility.Vector3dVector(pts)
    ls.lines = o3d.utility.Vector2iVector(lines)
    ls.colors = o3d.utility.Vector3dVector(colors)
    return ls


def create_point_cloud(points, size=5.0):
    pts = np.asarray([[p[0], p[1], p[2]] for p in points], dtype=float)
    pc = o3d.geometry.PointCloud()
    pc.points = o3d.utility.Vector3dVector(pts)
    colors = gradient_colors(len(points))
    pc.colors = o3d.utility.Vector3dVector(colors)
    return pc


def visualize(inner_points, outer_points):
    vis = o3d.visualization.Visualizer()
    vis.create_window(window_name='Gluing Path Viewer', width=1400, height=900)


    geometries = []
    # 添加坐标轴
    axis = o3d.geometry.TriangleMesh.create_coordinate_frame(size=20.0, origin=[0,0,0])
    geometries.append(axis)
    if inner_points:
        pc_inner = create_point_cloud(inner_points)
        geometries.append(pc_inner)
        ls_inner = make_line_set(inner_points)
        if ls_inner: geometries.append(ls_inner)
    if outer_points:
        pc_outer = create_point_cloud(outer_points)
        geometries.append(pc_outer)
        ls_outer = make_line_set(outer_points)
        if ls_outer: geometries.append(ls_outer)

    for g in geometries:
        vis.add_geometry(g)

    # 使用更好的渲染选项
    opt = vis.get_render_option()
    opt.point_size = 5.0
    opt.background_color = np.asarray([0.0, 0.0, 0.0])

    # 打印坐标信息并添加小球作为标签（Open3D 没有 3D text）
    if inner_points:
        print('\n=== 内部路径点坐标 ===')
        add_text_labels(vis, inner_points, 'Inner')
    if outer_points:
        print('\n=== 外部路径点坐标 ===')
        add_text_labels(vis, outer_points, 'Outer')

    vis.run()
    vis.destroy_window()


def make_line_set(points):
    """根据点创建线集合，用灰色线连接相邻点"""
    if len(points) < 2:
        return None
    pts = np.asarray([[p[0], p[1], p[2]] for p in points], dtype=float)
    lines = [[i, i+1] for i in range(len(pts)-1)]
    colors = [[0.6, 0.6, 0.6] for _ in lines]
    ls = o3d.geometry.LineSet()
    ls.points = o3d.utility.Vector3dVector(pts)
    ls.lines = o3d.utility.Vector2iVector(lines)
    ls.colors = o3d.utility.Vector3dVector(colors)
    return ls


def create_point_cloud(points, size=5.0):
    pts = np.asarray([[p[0], p[1], p[2]] for p in points], dtype=float)
    pc = o3d.geometry.PointCloud()
    pc.points = o3d.utility.Vector3dVector(pts)
    colors = gradient_colors(len(points))
    pc.colors = o3d.utility.Vector3dVector(colors)
    return pc


def add_text_labels(vis, points, prefix):
    """在可视化窗口中为每个点显示局部坐标轴（使用点的 rx,ry,rz），并在终端打印信息。

    说明:
    - 假设 rx,ry,rz 单位为弧度，对应绕 X, Y, Z 的旋转（按 X then Y then Z 应用），
      计算旋转矩阵 R = Rz * Ry * Rx，然后将坐标轴旋转并移动到点位置。
    - Open3D 原生没有 3D 文本，这里用坐标轴网格显示局部坐标系。
    """
    def euler_to_rotation_matrix(rx, ry, rz):
        cx = np.cos(rx); sx = np.sin(rx)
        cy = np.cos(ry); sy = np.sin(ry)
        cz = np.cos(rz); sz = np.sin(rz)

        Rx = np.array([[1, 0, 0], [0, cx, -sx], [0, sx, cx]])
        Ry = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]])
        Rz = np.array([[cz, -sz, 0], [sz, cz, 0], [0, 0, 1]])

        # 应用顺序: 首先绕X, 然后Y, 然后Z -> 总变换 R = Rz * Ry * Rx
        return Rz.dot(Ry.dot(Rx))

    axis_size = 6.0
    for i, p in enumerate(points):
        x, y, z, rx, ry, rz, region = p
        # 小球标记点位
        sphere = o3d.geometry.TriangleMesh.create_sphere(radius=0.6)
        sphere.translate(np.array([x, y, z]))
        c = gradient_colors(len(points))[i]
        sphere.paint_uniform_color(c)
        sphere.compute_vertex_normals()
        vis.add_geometry(sphere)

        # 为每个点创建局部坐标系并根据姿态旋转
        frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=axis_size, origin=[0, 0, 0])
        R = euler_to_rotation_matrix(rx, ry, rz)
        frame.rotate(R, center=(0, 0, 0))
        frame.translate(np.array([x, y, z]))
        vis.add_geometry(frame)

        # 在终端显示坐标和姿态
        print(f"{prefix} 点 {i}: x={x:.2f}, y={y:.2f}, z={z:.2f}, rx={rx:.4f}, ry={ry:.4f}, rz={rz:.4f}, region={region}  --> 局部坐标轴已显示")


def add_text_labels_with_numbers(vis, points, prefix):
    """在可视化窗口中为每个点显示局部坐标轴和序号标记，并在终端打印信息。
    
    为了显示序号，我们在每个点附近放置一个大一点的球体来表示序号。
    通过球体的大小和颜色来区分不同的序号范围。
    """
    def euler_to_rotation_matrix(rx, ry, rz):
        cx = np.cos(rx); sx = np.sin(rx)
        cy = np.cos(ry); sy = np.sin(ry)
        cz = np.cos(rz); sz = np.sin(rz)

        Rx = np.array([[1, 0, 0], [0, cx, -sx], [0, sx, cx]])
        Ry = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]])
        Rz = np.array([[cz, -sz, 0], [sz, cz, 0], [0, 0, 1]])

        # 应用顺序: 首先绕X, 然后Y, 然后Z -> 总变换 R = Rz * Ry * Rx
        return Rz.dot(Ry.dot(Rx))

    axis_size = 4.0  # 稍微小一点的坐标轴，给序号标记留出空间
    
    for i, p in enumerate(points):
        x, y, z, rx, ry, rz, region = p
        
        # 主点球体（用渐变色表示顺序）
        sphere = o3d.geometry.TriangleMesh.create_sphere(radius=0.8)
        sphere.translate(np.array([x, y, z]))
        c = gradient_colors(len(points))[i]
        sphere.paint_uniform_color(c)
        sphere.compute_vertex_normals()
        vis.add_geometry(sphere)

        # 序号标记球体（白色，稍微偏移位置）
        number_sphere = o3d.geometry.TriangleMesh.create_sphere(radius=1.2)
        # 在Z轴上方偏移显示序号
        number_sphere.translate(np.array([x, y, z + 8.0]))
        number_sphere.paint_uniform_color([1.0, 1.0, 1.0])  # 白色
        number_sphere.compute_vertex_normals()
        vis.add_geometry(number_sphere)
        
        # 创建数字文本的替代方案：用不同颜色的小球组成数字的形状
        # 这里简化为用序号模10的结果来区分颜色
        digit_color = [
            [1.0, 0.0, 0.0],  # 0 - 红色
            [0.0, 1.0, 0.0],  # 1 - 绿色
            [0.0, 0.0, 1.0],  # 2 - 蓝色
            [1.0, 1.0, 0.0],  # 3 - 黄色
            [1.0, 0.0, 1.0],  # 4 - 品红色
            [0.0, 1.0, 1.0],  # 5 - 青色
            [1.0, 0.5, 0.0],  # 6 - 橙色
            [0.5, 0.0, 1.0],  # 7 - 紫色
            [0.0, 0.5, 0.0],  # 8 - 深绿色
            [0.5, 0.5, 0.5],  # 9 - 灰色
        ]
        digit_sphere = o3d.geometry.TriangleMesh.create_sphere(radius=0.5)
        digit_sphere.translate(np.array([x + 3.0, y, z + 8.0]))
        digit_sphere.paint_uniform_color(digit_color[i % 10])
        digit_sphere.compute_vertex_normals()
        vis.add_geometry(digit_sphere)

        # 为每个点创建局部坐标系并根据姿态旋转
        frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=axis_size, origin=[0, 0, 0])
        R = euler_to_rotation_matrix(rx, ry, rz)
        frame.rotate(R, center=(0, 0, 0))
        frame.translate(np.array([x, y, z]))
        vis.add_geometry(frame)

        # 在终端显示坐标和姿态（包含序号）
        print(f"{prefix} 点 [{i:03d}]: x={x:.2f}, y={y:.2f}, z={z:.2f}, rx={rx:.4f}, ry={ry:.4f}, rz={rz:.4f}, region={region}")


def visualize_global(global_points):
    """专门用于显示全局路径的可视化函数"""
    vis = o3d.visualization.Visualizer()
    vis.create_window(window_name='Global Sweep Path Viewer', width=1400, height=900)

    geometries = []
    
    # 添加全局坐标轴
    axis = o3d.geometry.TriangleMesh.create_coordinate_frame(size=30.0, origin=[0,0,0])
    geometries.append(axis)
    
    if global_points:
        # 创建基础点云（不使用颜色渐变）
        pts = np.asarray([[p[0], p[1], p[2]] for p in global_points], dtype=float)
        pc = o3d.geometry.PointCloud()
        pc.points = o3d.utility.Vector3dVector(pts)
        # 使用统一的蓝色
        pc.paint_uniform_color([0.0, 0.5, 1.0])
        geometries.append(pc)
        
        # 创建连线
        ls_global = make_line_set(global_points)
        if ls_global: 
            geometries.append(ls_global)

    # 添加所有几何体到可视化器
    for g in geometries:
        vis.add_geometry(g)

    # 设置渲染选项
    opt = vis.get_render_option()
    opt.point_size = 8.0
    opt.line_width = 2.0
    opt.background_color = np.asarray([0.1, 0.1, 0.1])

    # 为每个点添加局部坐标轴并输出终端信息
    if global_points:
        print('\n=== 全局扫掠路径点坐标 ===')
        print(f'总共 {len(global_points)} 个路径点')
        add_coordinate_axes(vis, global_points)

    # 运行可视化器
    vis.run()
    vis.destroy_window()


def add_coordinate_axes(vis, points):
    """为每个点添加局部坐标轴，根据rx,ry,rz旋转"""
    def euler_to_rotation_matrix(rx, ry, rz):
        cx = np.cos(rx); sx = np.sin(rx)
        cy = np.cos(ry); sy = np.sin(ry)
        cz = np.cos(rz); sz = np.sin(rz)

        Rx = np.array([[1, 0, 0], [0, cx, -sx], [0, sx, cx]])
        Ry = np.array([[cy, 0, sy], [0, 1, 0], [-sy, 0, cy]])
        Rz = np.array([[cz, -sz, 0], [sz, cz, 0], [0, 0, 1]])

        # 应用顺序: 首先绕X, 然后Y, 然后Z -> 总变换 R = Rz * Ry * Rx
        return Rz.dot(Ry.dot(Rx))

    axis_size = 6.0
    for i, p in enumerate(points):
        x, y, z, rx, ry, rz, region = p
        
        # 为每个点创建局部坐标系并根据姿态旋转
        frame = o3d.geometry.TriangleMesh.create_coordinate_frame(size=axis_size, origin=[0, 0, 0])
        R = euler_to_rotation_matrix(rx, ry, rz)
        frame.rotate(R, center=(0, 0, 0))
        frame.translate(np.array([x, y, z]))
        vis.add_geometry(frame)

        # 在终端显示坐标和姿态
        print(f"点 [{i:03d}]: x={x:.2f}, y={y:.2f}, z={z:.2f}, rx={rx:.4f}, ry={ry:.4f}, rz={rz:.4f}, region={region}")


def visualize(inner_points, outer_points):
    vis = o3d.visualization.Visualizer()
    vis.create_window(window_name='Gluing Path Viewer', width=1400, height=900)


    geometries = []
    # 添加坐标轴
    axis = o3d.geometry.TriangleMesh.create_coordinate_frame(size=20.0, origin=[0,0,0])
    geometries.append(axis)
    if inner_points:
        pc_inner = create_point_cloud(inner_points)
        geometries.append(pc_inner)
        ls_inner = make_line_set(inner_points)
        if ls_inner: geometries.append(ls_inner)
    if outer_points:
        pc_outer = create_point_cloud(outer_points)
        geometries.append(pc_outer)
        ls_outer = make_line_set(outer_points)
        if ls_outer: geometries.append(ls_outer)

    for g in geometries:
        vis.add_geometry(g)

    # 使用更好的渲染选项
    opt = vis.get_render_option()
    opt.point_size = 5.0
    opt.background_color = np.asarray([0.0, 0.0, 0.0])

    # 打印坐标信息并添加小球作为标签（Open3D 没有 3D text）
    if inner_points:
        print('\n=== 内部路径点坐标 ===')
        add_text_labels(vis, inner_points, 'Inner')
    if outer_points:
        print('\n=== 外部路径点坐标 ===')
        add_text_labels(vis, outer_points, 'Outer')

    vis.run()
    vis.destroy_window()


def main():
    parser = argparse.ArgumentParser(description='Global Gluing Path Visualizer (Open3D)')
    parser.add_argument('--global-path', type=str, help='全局路径文件路径', default=None)
    args = parser.parse_args()

    # 默认全局路径在 build/ 目录
    base_dir = Path(__file__).resolve().parent
    default_global = base_dir / 'build' / '_global_sweep_path.txt'

    global_file = Path(getattr(args, 'global_path')) if getattr(args, 'global_path') else default_global

    # 加载全局路径文件
    if global_file.exists():
        global_points = parse_path_file(str(global_file), is_outer=False)
        print(f"已加载全局路径文件: {global_file}")
        visualize_global(global_points)
    else:
        print(f"未找到全局路径文件: {global_file}")
        print("请确保文件存在或使用 --global-path 指定正确的路径")


if __name__ == '__main__':
    main()
